package cn.ymatrix.copyToYMatrix;

import cn.ymatrix.logger.LoggerLevel;
import cn.ymatrix.logger.MxLogger;
import cn.ymatrix.utils.Config;
import cn.ymatrix.utils.DataUtil;
import org.postgresql.copy.CopyManager;
import org.postgresql.core.BaseConnection;
import org.slf4j.Logger;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;

public class CopyStreamTOYMatrix {

  public static void main(String[] args) {

    MxLogger.loggerLevel(LoggerLevel.INFO);
    Logger LOGGER = MxLogger.init(CopyStreamTOYMatrix.class);

    String copySql =
        "COPY "
            + Config.getTableName()
            + " FROM STDIN DELIMITER E'"
            + Config.getDelimiter()
            + "' CSV";
    Connection connection = null;
    InputStream inputStream = null;
    int batchSize = Config.getBatchSize();
    int totalNumber = Config.getTotalNumber();
    try {
      connection =
          DriverManager.getConnection(Config.getUrl(), Config.getUserName(), Config.getPassWord());
      long startTime = System.currentTimeMillis();
      LOGGER.info("---- 开始插入数据 ----");
      CopyManager copyManager = new CopyManager((BaseConnection) connection);
      StringBuilder stringBuilder = new StringBuilder();
      for (int i = 1; i <= totalNumber; i++) {
        stringBuilder.append(
            DataUtil.getCurrentTimestamp()
                + "|"
                + DataUtil.getIntValue()
                + "|"
                + DataUtil.getIntValue()
                + "|"
                + DataUtil.getIntValue()
                + "|"
                + DataUtil.getIntValue()
                + "|"
                + DataUtil.getIntValue()
                + "\r\n");
        if (i % batchSize == 0) {
          inputStream =
              new ByteArrayInputStream(stringBuilder.toString().getBytes(StandardCharsets.UTF_8));
          copyManager.copyIn(copySql, inputStream);
          LOGGER.info("成功插入第 " + i + " 条数据");
          stringBuilder.delete(0, stringBuilder.length());
        }
      }
      // 提交末尾的数据
      inputStream =
          new ByteArrayInputStream(stringBuilder.toString().getBytes(StandardCharsets.UTF_8));
      copyManager.copyIn(copySql, inputStream);
      long spendTime = System.currentTimeMillis() - startTime;
      LOGGER.info("成功插入 " + totalNumber + " 条数据,耗时: " + spendTime + "毫秒");
      // stringBuilder.delete(0, stringBuilder.length());
    } catch (SQLException e) {
      LOGGER.error("SQLException error " + e.getMessage());
      e.printStackTrace();
    } catch (IOException e) {
      LOGGER.error("IOException error " + new RuntimeException(e));
    } finally {
      if (inputStream != null) {
        try {
          inputStream.close();
        } catch (IOException e) {
          LOGGER.error("inputStream close error " + e.getMessage());
        }
      }
      if (connection != null) {
        try {
          connection.close();
        } catch (SQLException e) {
          LOGGER.error("connection close error " + e.getMessage());
        }
      }
    }
  }
}
